import os
from launch import LaunchDescription
from launch_ros.actions import Node
from launch.actions import ExecuteProcess
from ament_index_python.packages import get_package_share_directory
from moveit_configs_utils import MoveItConfigsBuilder

def generate_launch_description():
    # 1️⃣ 修改 MoveIt 机器人配置
    moveit_config = (
        MoveItConfigsBuilder("zc4")  # ⚠️ 修改为你的 MoveIt 机器人包名
        .robot_description(file_path="config/zc4model.urdf.xacro")  # ⚠️ 替换 URDF 文件
        .robot_description_semantic(file_path="config/zc4model.srdf")  # ⚠️ 添加 SRDF
        .trajectory_execution(file_path="config/moveit_controllers.yaml")  # ⚠️ 替换控制器配置文件
        .to_moveit_configs()
    )

    #2️⃣ 启动 move_group 节点（MoveIt 核心）
    run_move_group_node = Node(
        package="moveit_ros_move_group",
        executable="move_group",
        output="screen",
        parameters=[moveit_config.to_dict()],
    )

    # 3️⃣ 启动 RViz
    rviz_config_file = (
        get_package_share_directory("move_group_interface") + "/launch/move_group.rviz"
    )  # ⚠️ 修改为你的 RViz 配置文件

    rviz_node = Node(
        package="rviz2",
        executable="rviz2",
        name="rviz2",
        output="log",
        arguments=["-d", rviz_config_file],
        parameters=[
            moveit_config.robot_description,
            moveit_config.robot_description_semantic,
            moveit_config.robot_description_kinematics,
            moveit_config.planning_pipelines,
            moveit_config.joint_limits,
        ],
    )

    # 4️⃣ 设置静态 TF（修改为你的机械臂基座）
    static_tf = Node(
        package="tf2_ros",
        executable="static_transform_publisher",
        name="static_transform_publisher",
        output="log",
        arguments=["--frame-id", "world", "--child-frame-id", "base_link"],  # ⚠️ 替换基座坐标系
    )

    # 5️⃣ 运行 robot_state_publisher（发布关节状态）
    robot_state_publisher = Node(
        package="robot_state_publisher",
        executable="robot_state_publisher",
        name="robot_state_publisher",
        output="both",
        parameters=[moveit_config.robot_description],
    )

    # 6️⃣ 设置 ros2_control 控制器配置文件
    ros2_controllers_path = os.path.join(
        get_package_share_directory("zc4_moveit_config"),  # ⚠️ 替换 MoveIt 配置包
        "config",
        "ros2_controllers.yaml",  
    )

    # 7️⃣ 运行 ros2_control 控制器管理器
    ros2_control_node = Node(
        package="controller_manager",
        executable="ros2_control_node",
        parameters=[moveit_config.robot_description,ros2_controllers_path], 
        remappings=[
            ("/controller_manager/robot_description", "/robot_description"),
        ],
        output="both",
    )

    # 8️⃣ 启动控制器（确保名称与 ros2_controllers.yaml 匹配）
    load_controllers = []
    for controller in [
        "zc4_controller",  # ⚠️ 替换为你的机械臂控制器
        "joint_state_broadcaster",
    ]:
        load_controllers += [
            ExecuteProcess(
                cmd=["ros2 run controller_manager spawner {}".format(controller)],
                shell=True,
                output="screen",
            )
        ]

    return LaunchDescription(
        [
            rviz_node,
            static_tf,
            robot_state_publisher,
            run_move_group_node,
            ros2_control_node,
        ]
        + load_controllers
    )
